"""
统计套利策略(配对交易)(Zipline实现)
策略逻辑：
1. 选择高度相关的资产对
2. 计算价差和标准差
3. 当价差偏离均值时建仓
4. 价差回归时平仓
"""

from zipline.api import order_target_percent, record, symbol, get_datetime
from zipline.finance import commission, slippage
import numpy as np
import pandas as pd

def initialize(context):
    # 策略参数
    context.symbol1 = symbol('AAPL')
    context.symbol2 = symbol('MSFT')
    context.lookback = 30
    context.z_threshold = 2
    context.spread_history = []
    context.position = 0
    
    # 设置交易成本
    context.set_commission(commission.PerShare(cost=0.001, min_trade_cost=1))
    context.set_slippage(slippage.FixedSlippage(spread=0.01))
    
    # 初始化价格历史
    context.price_history1 = []
    context.price_history2 = []

def calculate_hedge_ratio(self, price1, price2):
    """使用协整关系计算对冲比例"""
    if len(price1) < 2 or len(price2) < 2:
        return 1.0
    return np.cov(price1, price2)[0,1] / np.var(price2)

def calculate_spread(self, price1, price2):
    """计算标准化价差"""
    hedge_ratio = self.calculate_hedge_ratio(price1, price2)
    spread = price1 - hedge_ratio * price2
    return (spread - np.mean(spread)) / np.std(spread)

def handle_data(context, data):
    # 获取当前价格
    price1 = data.current(context.symbol1, 'price')
    price2 = data.current(context.symbol2, 'price')
    
    # 更新价格历史
    context.price_history1.append(price1)
    context.price_history2.append(price2)
    
    # 保持历史数据长度
    if len(context.price_history1) > context.lookback:
        context.price_history1.pop(0)
        context.price_history2.pop(0)
    
    # 检查是否有足够数据
    if len(context.price_history1) < context.lookback:
        return
    
    # 计算当前价差
    current_spread = calculate_spread(context, 
                                     np.array(context.price_history1),
                                     np.array(context.price_history2))[-1]
    
    # 生成信号
    signal = 0
    if current_spread > context.z_threshold and context.position <= 0:
        signal = -1  # 做空价差(买入symbol2，卖出symbol1)
    elif current_spread < -context.z_threshold and context.position >= 0:
        signal = 1   # 做多价差(买入symbol1，卖出symbol2)
    elif abs(current_spread) < 0.5 and context.position != 0:
        signal = 0    # 平仓
    
    # 执行交易
    if signal == 1:
        # 做多价差
        order_target_percent(context.symbol1, 0.5)
        order_target_percent(context.symbol2, -0.5)
        context.position = 1
    elif signal == -1:
        # 做空价差
        order_target_percent(context.symbol1, -0.5)
        order_target_percent(context.symbol2, 0.5)
        context.position = -1
    elif signal == 0:
        # 平仓
        order_target_percent(context.symbol1, 0)
        order_target_percent(context.symbol2, 0)
        context.position = 0
    
    # 记录状态
    record(price1=price1, 
          price2=price2, 
          spread=current_spread, 
          position=context.position)